# ocr_server.py
from flask import Flask, request, jsonify
import ddddocr
from PIL import Image, ImageEnhance, ImageFilter
import io
import traceback
import cv2
import numpy as np
import random
from paddleocr import PaddleOCR

app = Flask(__name__)
# 创建多个OCR实例用于集成识别
ocr_det = ddddocr.DdddOcr(det=True, show_ad=False)
ocr_cls = ddddocr.DdddOcr(det=False, show_ad=False)
ocr_beta = ddddocr.DdddOcr(det=False, beta=True, show_ad=False)  # 使用beta模型

# 初始化PaddleOCR（只需一次，建议全局）
paddle_ocr = PaddleOCR(use_angle_cls=True, lang='ch')


def preprocess_image_advanced(pil_img):
    """
    高级图像预处理函数，包含多种策略
    """
    try:
        # 转换为numpy数组
        img = np.array(pil_img)
        
        # 如果是RGB图像，转换为灰度图
        if len(img.shape) == 3:
            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        
        # 1. 图像增强 - 对比度调整
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        img = clahe.apply(img)
        
        # 2. 去噪 - 双边滤波（保持边缘）
        img = cv2.bilateralFilter(img, 9, 75, 75)
        
        # 3. 自适应二值化
        img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                  cv2.THRESH_BINARY, 15, 5)
        
        # 4. 形态学操作 - 去除小噪点
        kernel = np.ones((2, 2), np.uint8)
        img = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel)
        
        # 5. 边缘增强
        kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
        img = cv2.filter2D(img, -1, kernel)
        
        # 6. 再次去噪
        img = cv2.medianBlur(img, 3)
        
        # 转换回PIL图像
        return Image.fromarray(img)
        
    except Exception as e:
        print(f"高级预处理错误: {str(e)}")
        return pil_img
        


def preprocess_image_alternative(pil_img):
    """
    替代预处理策略，使用不同的参数
    """
    try:
        # 转换为numpy数组
        img = np.array(pil_img)
        
        # 如果是RGB图像，转换为灰度图
        if len(img.shape) == 3:
            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        
        # 1. 直方图均衡化
        img = cv2.equalizeHist(img)
        
        # 2. 高斯模糊
        img = cv2.GaussianBlur(img, (5, 5), 0)
        
        # 3. 自适应二值化（不同参数）
        img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_MEAN_C, 
                                  cv2.THRESH_BINARY, 21, 10)
        
        # 4. 形态学操作
        kernel = np.ones((3, 3), np.uint8)
        img = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
        
        # 转换回PIL图像
        return Image.fromarray(img)
        
    except Exception as e:
        print(f"替代预处理错误: {str(e)}")
        return pil_img


def enhance_image_size(pil_img, target_width=250):
    """
    调整图像大小，保持宽高比
    """
    try:
        width, height = pil_img.size
        if width < target_width:
            # 放大图像
            ratio = target_width / width
            new_width = int(width * ratio)
            new_height = int(height * ratio)
            return pil_img.resize((new_width, new_height), Image.LANCZOS)
        return pil_img
    except Exception as e:
        print(f"图像尺寸调整错误: {str(e)}")
        return pil_img


def rotate_image(pil_img):
    """
    尝试不同角度的旋转，找到最佳识别效果
    """
    results = []
    angles = [0, 90, 180, 270]  # 尝试0°, 90°, 180°, 270°
    
    for angle in angles:
        try:
            rotated = pil_img.rotate(angle, expand=True)
            results.append((rotated, angle))
        except Exception as e:
            print(f"旋转角度 {angle} 错误: {str(e)}")
    
    return results


def ensemble_ocr_recognition(processed_img):
    """
    集成OCR识别，使用多个模型和预处理方法
    """
    results = []
    
    try:
        # 方法1: 标准预处理 + 标准模型
        text1 = ocr_cls.classification(processed_img)
        results.append(text1)
        
        # 方法2: 标准预处理 + Beta模型
        text2 = ocr_beta.classification(processed_img)
        results.append(text2)
        
        # 方法3: 替代预处理 + 标准模型
        alt_img = preprocess_image_alternative(processed_img)
        text3 = ocr_cls.classification(alt_img)
        results.append(text3)
        
        # 方法4: 替代预处理 + Beta模型
        text4 = ocr_beta.classification(alt_img)
        results.append(text4)
        
        # 方法5: 旋转尝试
        rotated_results = rotate_image(processed_img)
        for rotated_img, angle in rotated_results:
            try:
                text_rot = ocr_cls.classification(rotated_img)
                if text_rot and text_rot.strip():
                    results.append(text_rot)
            except Exception as e:
                print(f"旋转识别错误: {str(e)}")
        
        # 投票机制：选择出现次数最多的结果
        if results:
            # 过滤空结果
            valid_results = [r for r in results if r and r.strip()]
            if valid_results:
                # 简单投票：选择第一个非空结果
                return valid_results[0]
        
        return results[0] if results else ""
        
    except Exception as e:
        print(f"集成OCR识别错误: {str(e)}")
        return ""


@app.route('/', methods=['GET', 'POST'])
def root():
    try:
        if request.method == 'GET':
            return jsonify({"message": "OCR Server is running", "endpoints": ["/ocr"]})
        else:
            # 如果是POST请求，重定向到OCR处理
            return ocr_image()
    except Exception as e:
        print(f"Error in root handler: {str(e)}")
        print(traceback.format_exc())
        return jsonify({"error": str(e)}), 500


@app.route('/ocr', methods=['POST'])
def ocr_image():
    try:
        # 接收图片
        if 'image' not in request.files:
            return jsonify({"error": "No image file provided"}), 400
        
        file = request.files['image']
        if file.filename == '':
            return jsonify({"error": "No file selected"}), 400
        
        print(f"Received file: {file.filename}")
        
        # 先读取文件字节，用于detection
        file_bytes = file.read()
        det_res = ocr_det.detection(file_bytes)
        print(f"Detection results: {det_res}")
        
        # 再用PIL打开图片，用于裁剪和classification
        img = Image.open(io.BytesIO(file_bytes))
        print(f"Image size: {img.size}")
        
        rec_res = []
        for i, box in enumerate(det_res):
            try:
                # 裁剪图像
                cropped_img = img.crop(box)
                
                # 图像预处理
                processed_img = preprocess_image_advanced(cropped_img)
                
                # 调整图像大小
                processed_img = enhance_image_size(processed_img)
                
                # 集成OCR识别
                text = ensemble_ocr_recognition(processed_img)
                rec_res.append(text)
                print(f"Box {i}: {text}")
                
            except Exception as e:
                print(f"Error processing box {box}: {str(e)}")
                rec_res.append("")

        # 构建返回数据
        result = []
        for (x1, y1, x2, y2), text in zip(det_res, rec_res):
            # 计算中心坐标
            center_x = (x1 + x2) // 2
            center_y = (y1 + y2) // 2
            result.append({
                "text": text,
                "x": int(center_x),
                "y": int(center_y),
                "width": x2 - x1,
                "height": y2 - y1,
                "center": {
                    "x": int(center_x),
                    "y": int(center_y)
                }
            })

        return jsonify({"code": 0, "data": result})
        
    except Exception as e:
        print(f"Error in OCR processing: {str(e)}")
        print(traceback.format_exc())
        return jsonify({"error": str(e)}), 500


@app.route('/ocr_paddle', methods=['POST'])
def ocr_paddle_image():
    try:
        if 'image' not in request.files:
            return jsonify({"error": "No image file provided"}), 400

        file = request.files['image']
        if file.filename == '':
            return jsonify({"error": "No file selected"}), 400

        file_bytes = file.read()
        img_path = 'temp_paddleocr_img.png'
        with open(img_path, 'wb') as f:
            f.write(file_bytes)

        # PaddleOCR 识别
        result = paddle_ocr.ocr(img_path, cls=True)
        data = []
        for line in result:
            for box, (text, conf) in line:
                # box: 4点坐标，text: 识别内容，conf: 置信度
                x_list = [int(pt[0]) for pt in box]
                y_list = [int(pt[1]) for pt in box]
                center_x = sum(x_list) // 4
                center_y = sum(y_list) // 4
                data.append({
                    "text": text,
                    "confidence": float(conf),
                    "box": box,
                    "center": {"x": center_x, "y": center_y}
                })

        return jsonify({"code": 0, "data": data})

    except Exception as e:
        print(f"Error in PaddleOCR processing: {str(e)}")
        import traceback
        print(traceback.format_exc())
        return jsonify({"error": str(e)}), 500


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001)